# -*- coding: utf-8 -*-

import logger.logger as logger
from sklearn.externals import joblib
from sklearn.ensemble import RandomForestClassifier



class RandomForest():
    def __init__(self, n_estimators=10,
                 criterion="entropy",
                 max_depth=None,
                 min_samples_split=2,
                 min_samples_leaf=1,
                 min_weight_fraction_leaf=0.,
                 max_features="auto",
                 max_leaf_nodes=None,
                 min_impurity_decrease=0.,
                 min_impurity_split=None,
                 bootstrap=True,
                 oob_score=False,
                 n_jobs=1,
                 random_state=None,
                 verbose=0,
                 warm_start=False,
                 class_weight=None):
        self.model = None
        self.parameter = {}
        self.parameter['n_estimators'] = n_estimators
        self.parameter['criterion'] = criterion
        self.parameter['max_depth'] = max_depth
        self.parameter['min_samples_split'] = min_samples_split
        self.parameter['min_samples_leaf'] = min_samples_leaf
        self.parameter['min_weight_fraction_leaf'] = min_weight_fraction_leaf
        self.parameter['max_features'] = max_features
        self.parameter['max_leaf_nodes'] = max_leaf_nodes
        self.parameter['min_impurity_decrease'] = min_impurity_decrease
        self.parameter['bootstrap'] = bootstrap
        self.parameter['oob_score'] = oob_score
        self.parameter['n_jobs'] = n_jobs
        self.parameter['random_state'] = random_state
        self.parameter['verbose'] = verbose
        self.parameter['warm_start'] = warm_start
        self.parameter['class_weight'] = class_weight


    def train(self, x_train, y_train, x_test, y_test, model_path):
        logger.info("training...")
        self.model = RandomForestClassifier(**self.parameter)
        self.model.fit(x_train, y_train)
        joblib.dump(self.model, model_path)
        acc = self.model.score(x_test, y_test)
        logger.info("acc: {}".format(str(acc)))

    def cross_validation(self, train_x, train_y, test_vecs, y_test, model_path, param_grid=None):
        from sklearn.model_selection import GridSearchCV
        model = RandomForestClassifier()
        param_grid = [
            {'kernel': ['rbf'],
             'C': [0.001, 0.01, 0.1, 1, 10, 100],
             'gamma': [0.001, 0.01, 0.1, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4, 1.6]},

            {'kernel': ['linear'],
             'C': [0.001, 0.01, 0.1, 1, 10, 10]}
        ]
        grid_search = GridSearchCV(model, param_grid, n_jobs=4, verbose=1, cv=5)
        grid_search.fit(train_x, train_y)
        best_parameters = grid_search.best_estimator_.get_params()

        for para, val in list(best_parameters.items()):
            print(para, val)

        self.model = RandomForestClassifier(kernel='rbf', C=best_parameters['C'], gamma=best_parameters['gamma'],
                                            probability=True)
        self.model.fit(train_x, train_y)
        joblib.dump(self.model, model_path)
        acc = self.model.score(test_vecs, y_test)
        logger.info("acc: {}".format(str(acc)))

    def load_model(self, model_path):
        self.model = joblib.load(model_path)

    def predict(self, x_vecs):
        if not self.model:
            logger.info("please load random forest model first!!!")
            return
        y_pred = self.model.predict(x_vecs)
        return y_pred


if __name__ == '__main__':
    pass